#! /usr/bin/env python3
import keras
from keras.models import load_model

import pickle
from glob import glob
import os

from tqdm import tqdm
from PIL import Image
import numpy as np

import fire

def main(model_filename, image_dirname, label_filename=None, class_name=None, output=None, input_shape=(128, 128, 3)):
    model = load_model(model_filename)

    if label_filename is not None:
        with open(label_filename, 'rb') as f:
            label = pickle.load(f)

    image_filenames = glob(os.path.join(image_dirname, '*'))

    def formatImage(image, input_shape):
        src_size = np.array(image.size)
        dst_size = np.array([input_shape[0], input_shape[1]])

        if image.mode != 'RGB': image = image.convert('RGB')
        trans_array = np.zeros(input_shape, dtype='uint8') + 255

        k = dst_size / src_size.max()
        scale_size = np.around((k * src_size)).astype('int')
        trans_vect = ((dst_size - scale_size) / 2).astype('int')

        scale_image = image.resize(tuple(scale_size), resample=Image.BOX)
        scale_array = np.array(scale_image, dtype='uint8')
        trans_array[trans_vect[1]: trans_vect[1]+scale_size[1],
                    trans_vect[0]: trans_vect[0]+scale_size[0]] =\
                    np.array(scale_image, dtype='uint8')

        return trans_array

    with open('error_image.csv', 'w') as f:
        for image_filename in tqdm(image_filenames):
            img = Image.open(image_filename).convert(('L', 'RGB')[input_shape[-1] == 3])
            arr = formatImage(img, input_shape)

            input_X = 1 - arr.astype('float32') / 255
            input_X = input_X.reshape((-1,) + input_shape)

            probs = model.predict(input_X)[0]
            predict = (probs == probs.max()).dot(np.arange(probs.size))
            if label_filename is not None:
                predict_class = label[predict]
            else:
                predict_class = str(predict)

            if predict_class != class_name:
                f.writelines([image_filename, ',', predict_class, '\n'])
                if output is not None:
                    output_dir = os.path.join(output, predict_class)
                    if not os.path.isdir(output_dir): os.mkdir(output_dir)
                    output_path = os.path.join(output_dir, os.path.basename(image_filename))
                    img.save(output_path)

if __name__ == '__main__':
    fire.Fire(main)
